# 导入所需的库
from zhipuai import ZhipuAI
import logging
import base64

# 配置日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

class Glm:
    """
    Glm类用于初始化模型并提供图像到Base64编码及创建对话完成的功能。
    """
    
    def __init__(self, config):  
        """
        初始化Glm类。
        
        参数:
        - config: 包含模型和API密钥的配置字典。
        """
        logger.info("Glm init")
        self.model = config.get("model")
        self.api_key = config.get("api_key")
        self.client = ZhipuAI(api_key=self.api_key) # 初始化ZhipuAI客户端
    
    def img2base64(self, img_path):
        """
        将图像文件转换为Base64编码的字符串。
        
        参数:
        - img_path: 图像文件的路径。
        
        返回:
        - Base64编码的字符串。
        """
        with open(img_path, 'rb') as img_file:
            img_base = base64.b64encode(img_file.read()).decode('utf-8')
        return img_base

    def create(self, text, img_path):
        """
        创建对话完成，包含文本和图像信息。
        
        参数:
        - text: 用户的文本信息。
        - img_path: 图像文件的路径。
        
        返回:
        - 模型生成的对话完成内容。
        """
        response = self.client.chat.completions.create(
            model= self.model,  # 指定调用的模型名称
            messages=[
                { 
                    "role": "user", 
                    "content": 
                    [ 
                        { 
                            "type": "image_url", 
                            "image_url": 
                            { 
                                "url" : self.img2base64(img_path=img_path) 
                            } 
                        }, 
                        { 
                            "type": "text", 
                            "text": text
                        } 
                    ] 
                }
            ]
        )
        result = response.choices[0].message.content
        logger.info("Glm response: " + result)
        return result